package themis.flink.riskengine;

import com.alibaba.fastjson.JSONObject;
import java.util.List;
import java.util.Map;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.util.Collector;

public class FeatureReduceFunction extends RichFlatMapFunction<JSONObject, JSONObject> {
    private List<Rule> features;

    private ValueState<JSONObject> merged;

    public FeatureReduceFunction(List<Rule> features) {
        this.features = features;
    }

    public void open(Configuration config) {
        this.merged = getRuntimeContext().getState(new ValueStateDescriptor("saved reduceJson", JSONObject.class));
    }

    public void flatMap(JSONObject value, Collector<JSONObject> out) throws Exception {
        JSONObject mergedValue = (JSONObject)this.merged.value();
        if (mergedValue == null)
            mergedValue = new JSONObject();
        String keyName = value.getString("KEY_NAME");
        if (keyName.equals("event")) {
            mergedValue.put("event", value);
        } else {
            mergedValue.putIfAbsent("features", new JSONObject());
            if (value.containsKey("features"))
                mergedValue.getJSONObject("features").putAll((Map)value.getJSONObject("features"));
        }
        if (mergedValue.containsKey("event") && mergedValue.containsKey("features") && mergedValue
                .getJSONObject("features").size() == this.features.size()) {
            out.collect(mergedValue);
            this.merged.clear();
        } else {
            this.merged.update(mergedValue);
        }
    }
}
